import onnx
import numpy as np
import onnx.helper as helper
from numpy.linalg import norm
import datetime

# 加载原始模型和量化模型
origin_path = "your_model.onnx"
quant_path = "your_model_quant.onnx"

origin_model = onnx.load(origin_path)
quant_model = onnx.load(quant_path)

def get_initializer_dict(graph):
    return {init.name: onnx.numpy_helper.to_array(init) for init in graph.initializer}

def get_quantization_params(graph):
    """提取量化参数（缩放因子和零点）"""
    quant_params = {}
    initializers = get_initializer_dict(graph)
    
    # 查找缩放因子和零点
    for name, value in initializers.items():
        if '_scale' in name.lower() or name.endswith('_scale'):
            # 提取对应的权重名称
            weight_name = name.replace('_scale', '').replace('_Scale', '')
            if weight_name not in quant_params:
                quant_params[weight_name] = {}
            quant_params[weight_name]['scale'] = value
            
        elif '_zero_point' in name.lower() or name.endswith('_zero_point'):
            # 提取对应的权重名称
            weight_name = name.replace('_zero_point', '').replace('_ZeroPoint', '')
            if weight_name not in quant_params:
                quant_params[weight_name] = {}
            quant_params[weight_name]['zero_point'] = value
    
    return quant_params

def dequantize_weight(quantized_weight, scale, zero_point):
    """反量化权重: (quantized_value - zero_point) * scale"""
    if scale is None or zero_point is None:
        print(f"警告: 缺少量化参数，跳过反量化")
        return quantized_weight.astype(np.float32)
    
    # 确保输入是正确的数据类型
    quantized_weight = quantized_weight.astype(np.float32)
    scale = float(scale) if np.isscalar(scale) else scale.astype(np.float32)
    zero_point = float(zero_point) if np.isscalar(zero_point) else zero_point.astype(np.float32)
    
    # 反量化公式: (quantized_value - zero_point) * scale
    dequantized = (quantized_weight - zero_point) * scale
    return dequantized

def strip_quant_suffix(name):
    # 去除常见的量化后缀/前缀
    for suffix in ["_quantized", "_dequantized", "_quant", "_dequant", "QuantizeLinear", "DequantizeLinear"]:
        name = name.replace(suffix, "")
    return name

def find_matching_quant_params(weight_name, quant_params):
    """为权重名称查找对应的量化参数"""
    # 直接匹配
    if weight_name in quant_params:
        return quant_params[weight_name]
    
    # 尝试去除后缀匹配
    stripped_name = strip_quant_suffix(weight_name)
    if stripped_name in quant_params:
        return quant_params[stripped_name]
    
    # 尝试模糊匹配
    for param_name in quant_params:
        if stripped_name in param_name or param_name in stripped_name:
            return quant_params[param_name]
    
    return None

origin_weights = get_initializer_dict(origin_model.graph)
quant_weights = get_initializer_dict(quant_model.graph)

# 获取量化参数
print("正在提取量化参数...")
quant_params = get_quantization_params(quant_model.graph)
print(f"找到 {len(quant_params)} 组量化参数")

# 打印量化参数信息（调试用）
for param_name, params in list(quant_params.items())[:5]:  # 只打印前5个
    scale = params.get('scale')
    zero_point = params.get('zero_point')
    print(f"参数 {param_name}: scale={scale}, zero_point={zero_point}")

# 建立量化权重的去后缀名到原始名的映射
quant_name_map = {}
for qname in quant_weights:
    stripped = strip_quant_suffix(qname)
    quant_name_map[stripped] = qname

# 创建输出文件
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
output_file = f"weight_comparison_dequantized_{timestamp}.txt"

with open(output_file, 'w', encoding='utf-8') as f:
    f.write(f"权重比较结果（反量化版本） - {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
    f.write("=" * 60 + "\n")
    f.write(f"原始模型: {origin_path}\n")
    f.write(f"量化模型: {quant_path}\n")
    f.write("=" * 60 + "\n\n")
    
    print(f"原始模型权重数: {len(origin_weights)}")
    print(f"量化模型权重数: {len(quant_weights)}")
    print(f"量化参数组数: {len(quant_params)}")
    f.write(f"原始模型权重数: {len(origin_weights)}\n")
    f.write(f"量化模型权重数: {len(quant_weights)}\n")
    f.write(f"量化参数组数: {len(quant_params)}\n\n")

    matched = 0
    not_matched = 0
    dequantized_count = 0
    low_similarity_weights = []  # 保存相似度小于0.9的权重名称

    for oname, origin_w in origin_weights.items():
        stripped = strip_quant_suffix(oname)
        if stripped in quant_name_map:
            qname = quant_name_map[stripped]
            quant_w = quant_weights[qname]
            
            # 查找对应的量化参数
            quant_param = find_matching_quant_params(qname, quant_params)
            
            if quant_w.shape != origin_w.shape:
                msg = f"[跳过] {oname} <-> {qname} 形状不一致: {origin_w.shape} vs {quant_w.shape}"
                print(msg)
                f.write(msg + "\n")
                continue
            
            # 尝试反量化
            if quant_param and 'scale' in quant_param and 'zero_point' in quant_param:
                scale = quant_param['scale']
                zero_point = quant_param['zero_point']
                dequantized_w = dequantize_weight(quant_w, scale, zero_point)
                dequantized_count += 1
                
                f.write(f"\n[反量化] {qname}: scale={scale}, zero_point={zero_point}\n")
            else:
                # 如果没有找到量化参数，直接转换为float32
                dequantized_w = quant_w.astype(np.float32)
                f.write(f"\n[直接转换] {qname}: 未找到量化参数，直接转换为float32\n")
            
            # 计算MSE和余弦相似度
            mse = np.mean((origin_w.astype(np.float32) - dequantized_w.astype(np.float32)) ** 2)
            
            flat_origin = origin_w.flatten().astype(np.float32)
            flat_dequant = dequantized_w.flatten().astype(np.float32)
            cos_sim = np.dot(flat_origin, flat_dequant) / (norm(flat_origin) * norm(flat_dequant) + 1e-8)
            
            # 保存张量信息到文件
            f.write(f"[张量详细信息] {oname} <-> {qname}\n")
            f.write(f"形状: {origin_w.shape}, 总元素数: {flat_origin.size}\n")
            f.write(f"原始权重前20个值: {flat_origin[:20].tolist()}\n")
            f.write(f"反量化权重前20个值: {flat_dequant[:20].tolist()}\n")
            f.write(f"原始权重统计: min={flat_origin.min():.6f}, max={flat_origin.max():.6f}, mean={flat_origin.mean():.6f}\n")
            f.write(f"反量化权重统计: min={flat_dequant.min():.6f}, max={flat_dequant.max():.6f}, mean={flat_dequant.mean():.6f}\n")
            f.write(f"MSE: {mse:.6f}, CosineSim: {cos_sim:.6f}\n")
            f.write("-" * 40 + "\n")
            
            # 检查相似度是否小于0.9
            if cos_sim < 0.9:
                low_similarity_weights.append(oname)
            
            msg = f"[匹配] {oname} <-> {qname}: CosineSim={cos_sim:.6f} {'(反量化)' if quant_param else '(直接转换)'}"
            print(msg)
            f.write(msg + "\n")
            matched += 1
        else:
            msg = f"[未匹配] {oname} 在量化模型中未找到对应权重"
            print(msg)
            f.write(msg + "\n")
            not_matched += 1

    # 输出相似度小于0.9的权重信息
    if low_similarity_weights:
        print(f"\n发现 {len(low_similarity_weights)} 个相似度小于0.9的权重:")
        f.write(f"\n发现 {len(low_similarity_weights)} 个相似度小于0.9的权重:\n")
        f.write("-" * 60 + "\n")
        for oname in low_similarity_weights:
            low_sim_msg = f"{oname}"
            print(low_sim_msg)
            f.write(low_sim_msg + "\n")
        f.write("-" * 60 + "\n")
    else:
        print("\n所有匹配的权重相似度都大于等于0.9")
        f.write("\n所有匹配的权重相似度都大于等于0.9\n")

    summary = f"\n共匹配到 {matched} 个权重，未匹配到 {not_matched} 个权重。"
    summary += f"\n其中 {dequantized_count} 个权重进行了反量化，{matched - dequantized_count} 个权重直接转换。"
    print(summary)
    f.write(summary + "\n")

print(f"\n结果已保存到文件: {output_file}")